;; <trie> ::= (<branch>)
;; <branch> ::= (() . ()) | (<elt> <branch>*)
;;                ^ denotes a terminating node
;;
;; so a trie is essentially an assoc list

;; The merge procedure merges a list of sequences into a trie

(define (make-trie) '())

(define terminator (cons '() '()))

(define (merge seqs)
  (cond ((null? seqs)
	 '())
	((member '() seqs)
	 (cons terminator
               (merge (filter pair? seqs))))
	(else (let* ((p (partition-by-head (caar seqs) seqs))
		     (tails (car p))
		     (rest (cdr p)))
		(cons `(,(caar seqs) . ,(merge tails))
		      (merge rest))))))

(define (partition-by-head head seqs)
  (cond ((null? seqs)
         terminator)
	(else (let* ((p (partition-by-head head (cdr seqs)))
 		     (l (car p))
 		     (r (cdr p)))
		(if (equal? head (caar seqs))
		    (cons (cons (cdar seqs) l) r)
		    (cons l (cons (car seqs) r)))))))


;; Listing the elements

(define (enumerate-trie trie)
  (concat-map enumerate-branch trie))

(define (enumerate-branch branch)
  (cond ((equal? branch terminator)
	 (list '()))
	(else
	 (map (lambda (tail)
		(cons (car branch) tail))
	      (enumerate-trie (cdr branch))))))


;; Extending a trie one element at a time

(define (prepare-argument s)
  (cond ((list? s) s)
        ((string? s) (string->list s))
        ((symbol? s) (string->list (symbol->string s)))
        (else (error "Can't use in trie:" s))))

(define (extend-trie t s)
  (extend-trie* t (prepare-argument s)))

(define (extend-trie* t s)
  (if (null? s)
      (add-terminator t)
      (add-association t (car s) (cdr s))))

(define (add-terminator t)
  (cond ((assoc '() t) => (lambda (entry) t))
        (else (cons (cons '() '()) t))))

(define (add-association t symbol rest)
  (cond ((assoc symbol t) =>
         (lambda (entry)
           (set-cdr! entry (extend-trie* (cdr entry) rest))
           t))
        (else (cons (cons symbol (extend-trie* (make-trie) rest)) t))))


;; Querying a trie data structure

(define (lookup-trie t s)
  (lookup-trie* t (prepare-argument s)))

(define (lookup-trie* t s)
  (if (null? s)
      t
      (cond ((assoc (car s) t) =>
             (lambda (entry)
               (lookup-trie* (cdr entry) (cdr s))))
            (else #f))))

(define (trie-contains? t s)
  (cond ((lookup-trie t s) =>
         (lambda (t*)
           (cond ((assoc '() t*) #t)
                 (else #f))))
        (else #f)))



;; Printing trie

(define (spaces n)
  (when (> n 0)
    (display " ")
    (spaces (- n 1))))

(define (print-trie l i)
  (for-each (lambda (b) (print-branch b i #t)) l))

(define (print-branch b i ?)
  (when ?
    (spaces i))
  (cond ((equal? b terminator)
	 (display ".")
	 (newline))
	(else
	 (let ((l (length (string->list (symbol->string (car b))))))
	   (display (car b))
	   (display " ")
	   (when (pair? (cdr b))
	     (print-branch (cadr b) (+ i l 1) #f)
	     (for-each (lambda (t) (print-branch t (+ i l 1) #t)) (cddr b)))))))

